import copy
from ..utils import BaseModuleBuilder
from .voc import VOCDataset
from .cityscapes import CityScapesDataset
from .MyRopeGlass import MyRopeGlass
from .MyGlass import MyGlass
from .MyRope import MyRope
from .MyRopeGlass_plusbigmodel import MyRopeGlassBigModel

'''DatasetBuilder'''
class DatasetBuilder(BaseModuleBuilder):
    REGISTERED_MODULES = {
        'VOCDataset': VOCDataset, "CityScapesDataset":CityScapesDataset, "MyRopeGlass":MyRopeGlass, "MyRopeGlassBigModel":MyRopeGlassBigModel,
        'MyGlass':MyGlass, "MyRope":MyRope
    }
    '''build'''
    def build(self, mode, logger_handle, dataset_cfg):
        dataset_cfg = copy.deepcopy(dataset_cfg)
        train_cfg, test_cfg = dataset_cfg.pop('train', {}), dataset_cfg.pop('test', {})
        dataset_cfg.update(train_cfg if mode == 'TRAIN' else test_cfg)
        module_cfg = {
            'mode': mode, 'logger_handle': logger_handle, 'dataset_cfg': dataset_cfg, 'type': dataset_cfg['type'],
        }
        return super().build(module_cfg)


'''BuildDataset'''
BuildDataset = DatasetBuilder().build
